/*
 *  Copyright (C) 2021 the original author or authors.
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package we.plugin.response.template;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.sun.xml.internal.messaging.saaj.util.ByteOutputStream;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import we.plugin.FizzPluginFilter;
import we.plugin.FizzPluginFilterChain;
import we.plugin.response.template.el.SpelTemplateContext;
import we.plugin.response.template.el.TemplateContext;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

/**
 * @author liuwenchao
 */
@Slf4j
@Component(ResponseTemplatePluginFilter.RESPONSE_TEMPLATE_PLUGIN_FILTER)
public class ResponseTemplatePluginFilter implements FizzPluginFilter {

    public static final String RESPONSE_TEMPLATE_PLUGIN_FILTER = "responseTemplatePlugin";

    private static final String STATUS_CODE = "statusCode";
    private static final String CONTENT_TYPE = "contentType";
    private static final String RESP_BODY = "respBody";
    private static final String RULE = "rule";
    private static final String TEMPLATE_KEY = "template_key";
    private static final String RESPONSE_STRING = "response";

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, Map<String, Object> config) {

        Integer statusCode = (Integer) config.get(STATUS_CODE);
        String contentType = (String) config.get(CONTENT_TYPE);
        String respBody = (String) config.get(RESP_BODY);
        String ruleStr = JSON.toJSONString(config.get(RULE));

        final ServerHttpResponse response = exchange.getResponse();
        try {
            final List<Rule> rules = JSONObject.parseArray(ruleStr, Rule.class);
            // 如果规则为空, 返回默认值
            if (!CollectionUtils.isEmpty(rules)) {
                exchange = buildResponse(exchange, response, contentType, respBody, rules, statusCode);
            }
        } catch (JSONException e) {
            log.error("json 解析异常: ", e);
            return FizzPluginFilterChain.next(exchange);
        }
        return FizzPluginFilterChain.next(exchange);
    }

    private ServerWebExchange buildResponse(ServerWebExchange exchange, ServerHttpResponse response, String contentType, String respBody, List<Rule> rules, Integer statusCode) {
        final DataBufferFactory dataBufferFactory = response.bufferFactory();
        ServerHttpResponseDecorator responseDecorator = new ServerHttpResponseDecorator(response) {

            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> flux = (Flux<? extends DataBuffer>) body;

                    return super.writeWith(flux.buffer().map(dataBuffers -> {
                        ByteOutputStream outputStream = new ByteOutputStream();
                        DataBuffer join = dataBufferFactory.join(dataBuffers);
                        byte[] content = new byte[join.readableByteCount()];
                        join.read(content);
                        DataBufferUtils.release(join);

                        for (Rule rule : rules) {
                            // template key
                            if (StringUtils.equals(rule.getTemplateKey(), exchange.getAttribute(TEMPLATE_KEY))) {
                                return getDataBuffer(outputStream, respBody, response, dataBufferFactory, statusCode);
                            }
                            // http status
                            final HttpStatus httpStatus = response.getStatusCode();
                            if (httpStatus != null && rule.getHttpStatus().compareTo(httpStatus.value()) == 0) {
                                return getDataBuffer(outputStream, respBody, response, dataBufferFactory, statusCode);
                            }
                            // SPEL
                            if (StringUtils.isNotBlank(rule.getSpel()) && spelHandler(rule.getSpel(), new String(content))) {
                                return getDataBuffer(outputStream, respBody, response, dataBufferFactory, statusCode);
                            }
                        }
                        response.getHeaders().setContentLength(content.length);
                        return dataBufferFactory.wrap(content);
                    }));
                }
                return super.writeWith(body);
            }
        };
        responseDecorator.getHeaders().setContentType(MediaType.parseMediaType(contentType));
        return exchange.mutate().response(responseDecorator).build();
    }

    private Boolean spelHandler(final String spel, String responseBody) {
        TemplateContext context = new SpelTemplateContext();
        context.setVariable(RESPONSE_STRING, JSON.parse(responseBody));
        return context.getSimpleValue(spel, boolean.class);
    }

    private DataBuffer getDataBuffer(ByteOutputStream outputStream, String respBody, ServerHttpResponse response, DataBufferFactory dataBufferFactory, Integer statusCode) {
        final byte[] bytes = respBody.getBytes(StandardCharsets.UTF_8);
        outputStream.write(bytes);
        response.getHeaders().setContentLength(outputStream.size());
        response.setStatusCode(HttpStatus.valueOf(statusCode));
        return dataBufferFactory.wrap(outputStream.getBytes());
    }


}
